#！/usr/bin/env python
#!coding=utf-8

from dataclasses import dataclass
import os 
import sys
from tkinter import Y


import matplotlib.pyplot as plt
import numpy as np

import jax.numpy as jnp
import jax
from datetime import date, datetime

def echo_cost_time(start_time, end_time, name):
    print(name + " ==============")
    print("start_time is ", start_time)
    #print(start_time)
    print("end_time is ", end_time)
    use_time = end_time - start_time
    print("use_time is ", use_time)
    total_second = use_time.seconds + use_time.microseconds* 1.0 / 1000
    print("cost time is ",total_second)

def sum_logistic(x): 
  return jnp.sum(1.0 / (1.0 + jnp.exp(-x))) 


def np_run():
    #a = np.arange(15).reshape(3, 5)
    #print(a)
    key = jax.random.PRNGKey(4)
    print(key)
    pykey = np.random.random()
    print(pykey)

    size = 5000 
    start_time = datetime.now()
    x = np.random.normal(size=(size, size)).astype(np.float32)
    p1_time = datetime.now()
    echo_cost_time(start_time, p1_time, "np_random")
    y = jax.random.normal(key, (size, size), dtype=jnp.float32)
    p2_time = datetime.now()
    echo_cost_time(p1_time, p2_time, "jnp_random")


    x_small = jnp.arange(3.) 
    print(x_small)
    derivative_fn = jax.grad(sum_logistic) 
    print(derivative_fn(x_small)) 

def run():
    print("jax world!")
    x_np = np.linspace(0, 10, 1000)
    y_np = 2 * np.sin(x_np) * np.cos(x_np)

    x_jnp = jnp.linspace(0, 10, 1000)
    y_jnp = 2 * jnp.sin(x_jnp) * jnp.cos(x_jnp)
    plt.plot(x_jnp, y_jnp)
    #plt.plot(x_np, y_np)
    #plt.show()

if __name__ == "__main__":
    #run()
    np_run()